using RAGApp.Api.Data;
using RAGApp.Api.Models;
using RAGApp.Api.Models.DTOs;
using RAGApp.Api.Services.Interfaces;
using Microsoft.EntityFrameworkCore;

namespace RAGApp.Api.Services;

public class DocumentService : IDocumentService
{
    private readonly ApplicationDbContext _context;
    private readonly IEmbeddingService _embeddingService;

    public DocumentService(ApplicationDbContext context, IEmbeddingService embeddingService)
    {
        _context = context;
        _embeddingService = embeddingService;
    }

    public async Task<DocumentResponseDto> UploadDocumentAsync(DocumentUploadDto dto)
    {
        var document = new Document
        {
            Id = Guid.NewGuid(),
            Title = dto.Title,
            Content = dto.Content,
            Source = dto.Source,
            CreatedAt = DateTime.UtcNow,
            UpdatedAt = DateTime.UtcNow
        };

        _context.Documents.Add(document);
        await _context.SaveChangesAsync();

        // 处理文档分块
        var chunks = await ChunkDocumentAsync(document);
        foreach (var chunk in chunks)
        {
            var embedding = await _embeddingService.GenerateEmbeddingAsync(chunk.Content);
            chunk.Embedding = new Pgvector.Vector(embedding);
            _context.DocumentChunks.Add(chunk);
        }
        await _context.SaveChangesAsync();

        return MapToDto(document);
    }

    public async Task<List<DocumentResponseDto>> GetAllDocumentsAsync()
    {
        var documents = await _context.Documents.ToListAsync();
        return documents.Select(MapToDto).ToList();
    }

    public async Task<DocumentResponseDto?> GetDocumentByIdAsync(Guid id)
    {
        var document = await _context.Documents.FindAsync(id);
        return document == null ? null : MapToDto(document);
    }

    public async Task<bool> DeleteDocumentAsync(Guid id)
    {
        var document = await _context.Documents.FindAsync(id);
        if (document == null)
            return false;

        _context.Documents.Remove(document);
        await _context.SaveChangesAsync();
        return true;
    }

    public Task<List<DocumentChunk>> ChunkDocumentAsync(Document document, int chunkSize = 1000, int overlap = 200)
    {
        var chunks = new List<DocumentChunk>();
        var content = document.Content;
        var sentences = content.Split(new[] { ". ", "! ", "? " }, StringSplitOptions.RemoveEmptyEntries);
        
        var currentChunk = new List<string>();
        var currentLength = 0;

        foreach (var sentence in sentences)
        {
            var sentenceLength = sentence.Length;
            
            if (currentLength + sentenceLength > chunkSize && currentChunk.Count > 0)
            {
                // 创建新的文档块
                var chunkText = string.Join(" ", currentChunk);
                chunks.Add(new DocumentChunk
                {
                    Id = Guid.NewGuid(),
                    DocumentId = document.Id,
                    Content = chunkText,
                    CreatedAt = DateTime.UtcNow
                });

                // 保留重叠部分
                var overlapWords = currentChunk
                    .SelectMany(s => s.Split(' '))
                    .TakeLast(overlap)
                    .ToList();
                
                currentChunk.Clear();
                currentChunk.Add(string.Join(" ", overlapWords));
                currentLength = overlapWords.Sum(w => w.Length) + overlapWords.Count - 1;
            }

            currentChunk.Add(sentence);
            currentLength += sentenceLength + 1; // +1 for space
        }

        // 添加最后一个块
        if (currentChunk.Count > 0)
        {
            chunks.Add(new DocumentChunk
            {
                Id = Guid.NewGuid(),
                DocumentId = document.Id,
                Content = string.Join(" ", currentChunk),
                CreatedAt = DateTime.UtcNow
            });
        }

        return Task.FromResult(chunks);
    }

    private static DocumentResponseDto MapToDto(Document document)
    {
        return new DocumentResponseDto
        {
            Id = document.Id,
            Title = document.Title,
            Content = document.Content,
            Source = document.Source,
            CreatedAt = document.CreatedAt,
            UpdatedAt = document.UpdatedAt
        };
    }
} 